# middleware/redis_cluster_client.py
"""
Redis集群客户端实现
支持Redis Cluster模式，消除单点故障风险
"""

import asyncio
import json
import time
from typing import Any, Optional, Dict, List, Union
from contextlib import asynccontextmanager
from loguru import logger

try:
    from redis.asyncio import Redis as AsyncRedis
    from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster
    from redis.asyncio.sentinel import Sentinel as AsyncSentinel
    from redis.exceptions import RedisError, ConnectionError, TimeoutError
    REDIS_AVAILABLE = True
except ImportError:
    logger.warning("Redis集群支持未安装，将使用单节点模式")
    REDIS_AVAILABLE = False

from utils.constants import RedisMode, ErrorCode
from utils.exceptions import RedisConnectionError, RedisOperationError


class RedisClusterClient:
    """Redis集群客户端"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化Redis集群客户端
        
        Args:
            config: Redis配置字典
                - mode: Redis模式 (single/cluster/sentinel)
                - hosts: 主机列表
                - port: 端口
                - password: 密码
                - db: 数据库编号
                - max_connections: 最大连接数
                - timeout: 超时时间
                - retry_attempts: 重试次数
                - retry_delay: 重试延迟
        """
        self.config = config
        self.mode = config.get("mode", RedisMode.SINGLE)
        self.hosts = config.get("hosts", ["localhost"])
        self.port = config.get("port", 6379)
        self.password = config.get("password", None)
        self.db = config.get("db", 0)
        self.max_connections = config.get("max_connections", 100)
        self.timeout = config.get("timeout", 5.0)
        self.retry_attempts = config.get("retry_attempts", 3)
        self.retry_delay = config.get("retry_delay", 1.0)
        
        self._client = None
        self._is_connected = False
        self._connection_stats = {
            "total_connections": 0,
            "failed_connections": 0,
            "successful_operations": 0,
            "failed_operations": 0
        }
    
    async def initialize(self) -> bool:
        """初始化Redis连接"""
        try:
            logger.info(f"初始化Redis集群客户端，模式: {self.mode}")
            
            if not REDIS_AVAILABLE:
                logger.error("Redis库不可用，无法初始化集群客户端")
                return False
            
            if self.mode == RedisMode.CLUSTER:
                return await self._init_cluster()
            elif self.mode == RedisMode.SENTINEL:
                return await self._init_sentinel()
            else:
                return await self._init_single()
                
        except Exception as e:
            logger.error(f"Redis集群初始化失败: {e}")
            self._connection_stats["failed_connections"] += 1
            raise RedisConnectionError(f"Redis初始化失败: {e}")
    
    async def _init_cluster(self) -> bool:
        """初始化Redis集群"""
        try:
            startup_nodes = [
                {"host": host, "port": self.port} 
                for host in self.hosts
            ]
            
            self._client = AsyncRedisCluster(
                startup_nodes=startup_nodes,
                password=self.password,
                skip_full_coverage_check=True,
                max_connections=self.max_connections,
                socket_connect_timeout=self.timeout,
                socket_timeout=self.timeout,
                decode_responses=True,
                retry_on_timeout=True,
                retry_on_error=[ConnectionError, TimeoutError],
                retry=self.retry_attempts
            )
            
            # 测试连接
            await self._client.ping()
            self._is_connected = True
            self._connection_stats["total_connections"] += 1
            
            logger.success(f"Redis集群连接成功，节点数: {len(startup_nodes)}")
            return True
            
        except Exception as e:
            logger.error(f"Redis集群连接失败: {e}")
            raise
    
    async def _init_sentinel(self) -> bool:
        """初始化Redis哨兵模式"""
        try:
            # 解析哨兵配置
            sentinels = [(host, self.port) for host in self.hosts]
            
            sentinel = AsyncSentinel(
                sentinels,
                socket_timeout=self.timeout,
                socket_connect_timeout=self.timeout,
                password=self.password,
                decode_responses=True
            )
            
            # 获取主节点连接
            self._client = sentinel.master_for(
                'mymaster',  # 主节点名称
                socket_timeout=self.timeout,
                socket_connect_timeout=self.timeout,
                decode_responses=True
            )
            
            # 测试连接
            await self._client.ping()
            self._is_connected = True
            self._connection_stats["total_connections"] += 1
            
            logger.success("Redis哨兵模式连接成功")
            return True
            
        except Exception as e:
            logger.error(f"Redis哨兵模式连接失败: {e}")
            raise
    
    async def _init_single(self) -> bool:
        """初始化单节点Redis"""
        try:
            self._client = AsyncRedis(
                host=self.hosts[0],
                port=self.port,
                password=self.password,
                db=self.db,
                max_connections=self.max_connections,
                socket_connect_timeout=self.timeout,
                socket_timeout=self.timeout,
                decode_responses=True,
                retry_on_timeout=True,
                retry_on_error=[ConnectionError, TimeoutError]
            )
            
            # 测试连接
            await self._client.ping()
            self._is_connected = True
            self._connection_stats["total_connections"] += 1
            
            logger.success(f"单节点Redis连接成功: {self.hosts[0]}:{self.port}")
            return True
            
        except Exception as e:
            logger.error(f"单节点Redis连接失败: {e}")
            raise
    
    async def close(self):
        """关闭Redis连接"""
        if self._client:
            try:
                await self._client.close()
                self._is_connected = False
                logger.info("Redis连接已关闭")
            except Exception as e:
                logger.error(f"关闭Redis连接失败: {e}")
    
    def is_connected(self) -> bool:
        """检查连接状态"""
        return self._is_connected and self._client is not None
    
    async def _execute_with_retry(self, operation, key: str = None):
        """带重试机制执行操作"""
        for attempt in range(self.retry_attempts):
            try:
                if not self.is_connected():
                    await self.initialize()
                
                result = await operation()
                self._connection_stats["successful_operations"] += 1
                return result
                
            except (RedisError, ConnectionError, TimeoutError) as e:
                self._connection_stats["failed_operations"] += 1
                
                if attempt < self.retry_attempts - 1:
                    logger.warning(f"Redis操作失败，重试 {attempt + 1}/{self.retry_attempts}: {e}")
                    await asyncio.sleep(self.retry_delay * (2 ** attempt))  # 指数退避
                else:
                    logger.error(f"Redis操作最终失败: {e}")
                    raise RedisOperationError(f"Redis操作失败: {e}")
    
    # 基本操作
    async def get(self, key: str) -> Optional[str]:
        """获取值"""
        async def operation():
            return await self._client.get(key)
        
        return await self._execute_with_retry(operation, key)
    
    async def set(self, key: str, value: Union[str, dict, list], 
                  expire: int = None) -> bool:
        """设置值"""
        async def operation():
            # 序列化复杂类型
            if isinstance(value, (dict, list)):
                value = json.dumps(value, ensure_ascii=False)
            
            if expire:
                return await self._client.setex(key, expire, value)
            else:
                return await self._client.set(key, value)
        
        return await self._execute_with_retry(operation, key)
    
    async def delete(self, key: str) -> int:
        """删除键"""
        async def operation():
            return await self._client.delete(key)
        
        return await self._execute_with_retry(operation, key)
    
    async def exists(self, key: str) -> bool:
        """检查键是否存在"""
        async def operation():
            return bool(await self._client.exists(key))
        
        return await self._execute_with_retry(operation, key)
    
    async def expire(self, key: str, seconds: int) -> bool:
        """设置过期时间"""
        async def operation():
            return await self._client.expire(key, seconds)
        
        return await self._execute_with_retry(operation, key)
    
    # 哈希操作
    async def hget(self, name: str, key: str) -> Optional[str]:
        """获取哈希字段值"""
        async def operation():
            return await self._client.hget(name, key)
        
        return await self._execute_with_retry(operation, f"{name}:{key}")
    
    async def hset(self, name: str, key: str, value: Union[str, dict, list]) -> int:
        """设置哈希字段值"""
        async def operation():
            if isinstance(value, (dict, list)):
                value = json.dumps(value, ensure_ascii=False)
            return await self._client.hset(name, key, value)
        
        return await self._execute_with_retry(operation, f"{name}:{key}")
    
    async def hdel(self, name: str, key: str) -> int:
        """删除哈希字段"""
        async def operation():
            return await self._client.hdel(name, key)
        
        return await self._execute_with_retry(operation, f"{name}:{key}")
    
    async def hgetall(self, name: str) -> Dict[str, str]:
        """获取所有哈希字段"""
        async def operation():
            return await self._client.hgetall(name)
        
        return await self._execute_with_retry(operation, name)
    
    async def hkeys(self, name: str) -> List[str]:
        """获取哈希所有键"""
        async def operation():
            return await self._client.hkeys(name)
        
        return await self._execute_with_retry(operation, name)
    
    # 列表操作
    async def lpush(self, name: str, *values) -> int:
        """左推入列表"""
        async def operation():
            return await self._client.lpush(name, *values)
        
        return await self._execute_with_retry(operation, name)
    
    async def rpush(self, name: str, *values) -> int:
        """右推入列表"""
        async def operation():
            return await self._client.rpush(name, *values)
        
        return await self._execute_with_retry(operation, name)
    
    async def lpop(self, name: str) -> Optional[str]:
        """左弹出列表"""
        async def operation():
            return await self._client.lpop(name)
        
        return await self._execute_with_retry(operation, name)
    
    async def rpop(self, name: str) -> Optional[str]:
        """右弹出列表"""
        async def operation():
            return await self._client.rpop(name)
        
        return await self._execute_with_retry(operation, name)
    
    async def llen(self, name: str) -> int:
        """获取列表长度"""
        async def operation():
            return await self._client.llen(name)
        
        return await self._execute_with_retry(operation, name)
    
    # 集合操作
    async def sadd(self, name: str, *values) -> int:
        """添加集合成员"""
        async def operation():
            return await self._client.sadd(name, *values)
        
        return await self._execute_with_retry(operation, name)
    
    async def srem(self, name: str, *values) -> int:
        """移除集合成员"""
        async def operation():
            return await self._client.srem(name, *values)
        
        return await self._execute_with_retry(operation, name)
    
    async def smembers(self, name: str) -> set:
        """获取集合所有成员"""
        async def operation():
            return await self._client.smembers(name)
        
        return await self._execute_with_retry(operation, name)
    
    async def scard(self, name: str) -> int:
        """获取集合基数"""
        async def operation():
            return await self._client.scard(name)
        
        return await self._execute_with_retry(operation, name)
    
    # 发布订阅
    async def publish(self, channel: str, message: Union[str, dict, list]) -> int:
        """发布消息"""
        async def operation():
            if isinstance(message, (dict, list)):
                message = json.dumps(message, ensure_ascii=False)
            return await self._client.publish(channel, message)
        
        return await self._execute_with_retry(operation, channel)
    
    async def subscribe(self, *channels):
        """订阅频道"""
        pubsub = self._client.pubsub()
        await pubsub.subscribe(*channels)
        return pubsub
    
    # 分布式锁
    async def acquire_lock(self, lock_key: str, timeout: int = 10, 
                          blocking_timeout: int = 5) -> bool:
        """获取分布式锁"""
        try:
            lock_script = """
            if redis.call('exists', KEYS[1]) == 0 then
                redis.call('setex', KEYS[1], ARGV[2], ARGV[1])
                return 1
            elseif redis.call('get', KEYS[1]) == ARGV[1] then
                redis.call('expire', KEYS[1], ARGV[2])
                return 1
            else
                return 0
            end
            """
            
            lock_value = f"lock_{time.time()}_{asyncio.current_task().get_name()}"
            
            result = await self._client.eval(
                lock_script, 1, lock_key, lock_value, timeout
            )
            
            return bool(result)
            
        except Exception as e:
            logger.error(f"获取分布式锁失败: {e}")
            return False
    
    async def release_lock(self, lock_key: str) -> bool:
        """释放分布式锁"""
        try:
            unlock_script = """
            if redis.call('get', KEYS[1]) == ARGV[1] then
                return redis.call('del', KEYS[1])
            else
                return 0
            end
            """
            
            result = await self._client.eval(unlock_script, 1, lock_key, "lock_value")
            return bool(result)
            
        except Exception as e:
            logger.error(f"释放分布式锁失败: {e}")
            return False
    
    # 集群信息
    async def get_cluster_info(self) -> Dict[str, Any]:
        """获取集群信息"""
        try:
            if self.mode == RedisMode.CLUSTER:
                info = await self._client.cluster_info()
                nodes = await self._client.cluster_nodes()
                return {
                    "mode": "cluster",
                    "cluster_info": info,
                    "nodes": nodes
                }
            else:
                info = await self._client.info()
                return {
                    "mode": self.mode,
                    "info": info
                }
        except Exception as e:
            logger.error(f"获取集群信息失败: {e}")
            return {"mode": self.mode, "error": str(e)}
    
    def get_connection_stats(self) -> Dict[str, Any]:
        """获取连接统计"""
        return {
            **self._connection_stats,
            "is_connected": self.is_connected(),
            "mode": self.mode,
            "hosts": self.hosts,
            "port": self.port
        }


# 全局Redis集群客户端实例
_redis_cluster_client: Optional[RedisClusterClient] = None


async def get_redis_cluster_client(config: Dict[str, Any] = None) -> RedisClusterClient:
    """获取Redis集群客户端实例"""
    global _redis_cluster_client
    
    if _redis_cluster_client is None and config:
        _redis_cluster_client = RedisClusterClient(config)
        await _redis_cluster_client.initialize()
    
    return _redis_cluster_client